package com.jeffery.gateway.filter;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.annotation.Order;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Service;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Order(5)
@Service
@Slf4j
public class TenantCodeGatewayFilter implements GlobalFilter {

    private static final String TENANT_CODE_HEADER = "tenantcode";
    private static final String TENANT_CODE_PARAM = "tenantCode";

    @Resource
    private ObjectMapper objectMapper;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        ServerHttpRequest request = exchange.getRequest();
        HttpHeaders headers = request.getHeaders();
        Long userId = Long.valueOf(headers.get("userId").get(0));
        String realTenantCode = "1111111";
        HttpHeaders newQueryHeader = new HttpHeaders();


        for (String key : headers.keySet()) {
            if (newQueryHeader.containsKey(key)) {
                continue;
            }
            if (key.equalsIgnoreCase(TENANT_CODE_HEADER)) {
                List<String> tenantCodes = headers.get(TENANT_CODE_HEADER);
                String tenantCodeOfParam = null;
                if (CollUtil.isNotEmpty(tenantCodes)) {
                    tenantCodeOfParam = tenantCodes.get(0);
                }
                if (checkValidTenantCode(userId, Long.valueOf(tenantCodeOfParam))) {
                    newQueryHeader.put(TENANT_CODE_PARAM, Collections.singletonList(tenantCodeOfParam));
                } else {
                    newQueryHeader.put(TENANT_CODE_PARAM, Collections.singletonList(realTenantCode));
                }
                continue;
            }
            newQueryHeader.put(key, Objects.requireNonNull(headers.get(key)));
        }

        MultiValueMap<String, String> queryParams = new LinkedMultiValueMap<>(request.getQueryParams());
        MultiValueMap<String, String> newQueryParams = new LinkedMultiValueMap<>();
        for (String key : queryParams.keySet()) {
            if (newQueryParams.containsKey(key)) {
                continue;
            }
            if (key.equalsIgnoreCase(TENANT_CODE_PARAM)) {
                List<String> tenantCodes = queryParams.get(TENANT_CODE_PARAM);
                String tenantCodeOfParam = null;
                if (CollUtil.isNotEmpty(tenantCodes)) {
                    tenantCodeOfParam = tenantCodes.get(0);
                }
                if (checkValidTenantCode(userId, Long.valueOf(tenantCodeOfParam))) {
                    newQueryParams.put(TENANT_CODE_PARAM, Collections.singletonList(tenantCodeOfParam));
                } else {
                    newQueryParams.put(TENANT_CODE_PARAM, Collections.singletonList(realTenantCode));
                }
                continue;
            }
            newQueryParams.put(key, queryParams.get(key));
        }

        // Replace tenantCode in body
        Mono<String> modifiedBody = extractRequestBody(exchange)
                .map(body -> replaceWithConstantTenantCode(body, realTenantCode));

        ServerHttpRequest modifiedRequest = request.mutate()
                .headers(HttpHeaders::clear)
                .headers(httpHeaders -> httpHeaders.addAll(newQueryHeader))
                .uri(UriComponentsBuilder.fromUri(request.getURI())
                        .replaceQueryParams(newQueryParams)
                        .build().toUri())
                .build();
        HttpHeaders modifiedRequestHeaders = modifiedRequest.getHeaders();
        modifiedRequestHeaders.forEach((key, value) ->
                log("TenantCodeGatewayFilter request header result: {} - {}", key, value));

        MultiValueMap<String, String> modifiedRequestQueryParams = modifiedRequest.getQueryParams();
        modifiedRequestQueryParams.forEach((key, value) ->
                log("TenantCodeGatewayFilter request params result: {} - {}", key, value));

        log("TenantCodeGatewayFilter handle end");
        ServerWebExchange newServerWebExchange = exchange.mutate()
                .request(new ServerHttpRequestDecorator(modifiedRequest) {
                    @Override
                    public Flux<DataBuffer> getBody() {
                        return modifiedBody.flatMapMany(body -> {
                            DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(body.getBytes(StandardCharsets.UTF_8));
                            return Flux.just(buffer);
                        });
                    }

                    @Override
                    public MultiValueMap<String, String> getQueryParams() {
                        log("ServerHttpRequestDecorator:newQueryParams={}", JSONUtil.toJsonStr(newQueryParams));
                        return newQueryParams;
                    }

                    @Override
                    public HttpHeaders getHeaders() {
                        log("ServerHttpRequestDecorator:newQueryHeader={}", JSONUtil.toJsonStr(newQueryHeader));
                        return newQueryHeader;
                    }

                    @Override
                    public URI getURI() {
                        URI uri = super.getURI();
                        String query = uri.getQuery();
                        log("query = {}", query);
                        String rawQuery = uri.getRawQuery();
                        log("rawQuery = {}", rawQuery);
                        return uri;
                    }
                })
                .build();
        ServerHttpRequest replacedRequest = newServerWebExchange.getRequest();
        HttpHeaders replacedRequestHeaders = replacedRequest.getHeaders();
        replacedRequestHeaders.forEach((key, value) ->
                log("replacedRequestHeaders headers: {} - {}", key, value));

        MultiValueMap<String, String> replacedRequestQueryParams = replacedRequest.getQueryParams();
        replacedRequestQueryParams.forEach((key, value) ->
                log("replacedRequestHeaders params: {} - {}", key, value));
        return chain.filter(newServerWebExchange);
    }

    private void log(String format, Object... arguments) {
        log.info(format, arguments);
    }

    private static void log(String msg) {
        log.info(msg);
    }

    private static Mono<Void> doNothing(ServerWebExchange exchange, GatewayFilterChain chain) {
        log("TenantCodeGatewayFilter doNothing");
        return chain.filter(exchange);
    }

    private boolean checkValidTenantCode(Long userId, Long tenantCode) {
        return userId == 1111111 && tenantCode == 1111111;
    }



    private Mono<String> extractRequestBody(ServerWebExchange exchange) {
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .map(dataBuffer -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    return new String(bytes, StandardCharsets.UTF_8);
                })
                .onErrorResume(e -> Mono.just("{}"));
    }

    private String replaceWithConstantTenantCode(String body, String tenantCode) {
        try {
            Map<String, Object> jsonMap = objectMapper.readValue(body, Map.class);
            if (jsonMap.containsKey(TENANT_CODE_PARAM)) {
                jsonMap.remove(TENANT_CODE_PARAM);
                jsonMap.put(TENANT_CODE_PARAM, tenantCode);
            }

            return objectMapper.writeValueAsString(jsonMap);
        } catch (JsonProcessingException e) {
            log.warn("Error processing request body: {}", e.getMessage());
            return body;
        }
    }
}